#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
import time

# 导入P4Runtime库
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

def writeIpv4Rules(p4info_helper, sw, dst_ip, prefix_len, dst_mac, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, prefix_len)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": port
        }
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed: {dst_ip}/{prefix_len} -> MAC: {dst_mac}, Port: {port}")

def main():
    p4info_file_path = 'build/source_routing.p4.p4info.txt'
    bmv2_file_path = 'build/source_routing.json'

    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        print("=== 安装IPv4转发规则 ===")

        # 连接到三个交换机
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0)
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1)
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2)

        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        print("设置转发流水线配置...")
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        print("流水线配置设置完成")

        # 等待交换机就绪
        time.sleep(2)

        # 安装s1的规则
        print("配置交换机 s1:")
        writeIpv4Rules(p4info_helper, s1, "10.0.1.1", 32, "08:00:00:00:01:11", 1)
        writeIpv4Rules(p4info_helper, s1, "10.0.2.2", 32, "08:00:00:00:02:22", 2)
        writeIpv4Rules(p4info_helper, s1, "10.0.3.3", 32, "08:00:00:00:03:33", 3)

        # 安装s2的规则
        print("配置交换机 s2:")
        writeIpv4Rules(p4info_helper, s2, "10.0.1.1", 32, "08:00:00:00:01:11", 2)
        writeIpv4Rules(p4info_helper, s2, "10.0.2.2", 32, "08:00:00:00:02:22", 1)
        writeIpv4Rules(p4info_helper, s2, "10.0.3.3", 32, "08:00:00:00:03:33", 3)

        # 安装s3的规则
        print("配置交换机 s3:")
        writeIpv4Rules(p4info_helper, s3, "10.0.1.1", 32, "08:00:00:00:01:11", 2)
        writeIpv4Rules(p4info_helper, s3, "10.0.2.2", 32, "08:00:00:00:02:22", 3)
        writeIpv4Rules(p4info_helper, s3, "10.0.3.3", 32, "08:00:00:00:03:33", 1)

        print("✓ IPv4转发规则安装完成!")

    except Exception as e:
        print(f"错误: {e}")
        import traceback
        traceback.print_exc()

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    main()
